// Captain Proto -- Protobuf Object-capability RPC
// Copyright (c) 2009 Kenton Varda, Google Inc., and contributors
// http://capnproto.googlecode.com
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//     * Neither the name of the Captain Proto project nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package capnproto;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.Socket;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicInteger;

import com.google.protobuf.ByteString;
import com.google.protobuf.CodedInputStream;
import com.google.protobuf.CodedOutputStream;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.MessageLite;
import com.google.protobuf.UninitializedMessageException;

/**
 * Implements the Captain Proto protocol over a raw socket or input/output
 * stream pair.
 */
public final class Connection {
  private final CodedInputStream input;
  private final CodedOutputStream output;

  private final ArrayList<Channel> channels =
    new ArrayList<Channel>();
  private final AtomicInteger outgoingCallIdCounter = new AtomicInteger(0);

  private final Map<Integer, DoneCallback> pendingCalls =
      new HashMap<Integer, DoneCallback>();

  private boolean isRunning = false;

  // =================================================================

  /**
   * Create a new connection on the given socket.  The connection will not start
   * handling incoming messages until you call {@link #runEventLoop()}.
   */
  public Connection(final Socket socket) throws IOException {
    this(socket.getInputStream(), socket.getOutputStream());
  }

  /**
   * Create a new connection using the given input and output streams.  The
   * connection will not start handling incoming messages until you call
   * {@link #runEventLoop()}.
   */
  public Connection(InputStream input, OutputStream output) {
    this.input = CodedInputStream.newInstance(input);
    this.output = CodedOutputStream.newInstance(output);
    this.channels.add(new NullChannel());
  }

  /**
   * Set the default service to be exported over the connection.  The service
   * will immediately be callable by the remote end as soon as the event loop
   * starts.
   */
  public Connection setDefaultService(Channel service) {
    checkNotRunning();
    this.channels.set(0, service);
    return this;
  }

  /**
   * Get a channel connected to the default service exported by the other side
   * of the connection.  You can then call
   * {@code MyServiceType.fromChannel(channel)} to construct a new instance
   * of MySerivceType representing the remote service.
   */
  public Channel getDefaultService() {
    return new ChannelImpl(0);
  }

  /**
   * Reads and handles messages from the connection.
   *
   * Assuming no exceptions are raised, this will not return until the end of
   * the stream has been reached.  One way to force an early termination is
   * to close the input stream, though this will typically cause an exception
   * to be thrown.  If reading from a socket, another option is to call
   * {@link java.net.Socket#shutdownInput()}, which causes a clean EOF on the
   * input stream, causing {@code runEventLoop()} to return normally.
   *
   * Under no circumstances will {@code runEventLoop()} attempt to close the
   * input or output streams.  This is the responsibility of the caller.
   *
   * @param executor An executor on which to execute calls on local services.
   */
  public void runEventLoop(ExecutorService executor) throws IOException {
    checkNotRunning();
    isRunning = true;
    while (readOneMessage(executor)) {}
  }

  /**
   * Calls {@link #runEventLoop(ExecutorService)} using a cached thread pool
   * as the executor.
   */
  public void runEventLoop() throws IOException {
    ExecutorService executor = Executors.newCachedThreadPool();
    try {
      runEventLoop(executor);
    } finally {
      executor.shutdown();
    }
  }

  /**
   * Thrown by {@link #runEventLoop()} when an erroneous message is received
   * that forces termination of the connection.
   */
  public class ProtocolErrorException extends IOException {
    private ProtocolErrorException(String string) {
      super("Peer committed protocol error:  " + string);
    }

    private static final long serialVersionUID = 1L;
  }

  // =================================================================

  private void checkNotRunning() {
    if (isRunning) {
      throw new IllegalStateException("Connection is already running.");
    }
  }

  private boolean readOneMessage(ExecutorService executor) throws IOException {
    input.resetSizeCounter();

    int tag = input.readTag();

    if (tag == 0) {
      return false;
    }

    int number = tag >>> 3;

    // Must be length-delimited wire type.
    if ((tag & 0x07) != 2) {
      if (!input.skipField(tag)) {
        throw new IOException("Unexpected end-group tag.");
      }

      sendError(Proto.MiscError.Type.UNKNOWN_FIELD,
                "Saw unknown non-sub-message tag: " + tag,
                number);
      return true;
    }

    int size = input.readRawVarint32();
    int oldLimit = input.pushLimit(size);

    switch (number) {
      case Proto.Stream.HEADER_FIELD_NUMBER:
        handleMessage(Proto.SessionHeader.parseFrom(input));
        input.checkLastTagWas(0);
        break;
      case Proto.Stream.CALL_REQUEST_FIELD_NUMBER:
        handleMessage(Proto.CallRequest.parseFrom(input), executor);
        input.checkLastTagWas(0);
        break;
      case Proto.Stream.CALL_RESPONSE_FIELD_NUMBER:
        handleMessage(Proto.CallResponse.parseFrom(input));
        input.checkLastTagWas(0);
        break;
      case Proto.Stream.CANCEL_CALL_FIELD_NUMBER:
        handleMessage(Proto.CancelCall.parseFrom(input));
        input.checkLastTagWas(0);
        break;
      case Proto.Stream.ERROR_FIELD_NUMBER:
        handleMessage(Proto.MiscError.parseFrom(input));
        input.checkLastTagWas(0);
        break;
      default:
        sendError(Proto.MiscError.Type.UNKNOWN_FIELD,
                  "Saw unknown field number: " + tag,
                  number);
        input.skipRawBytes(size);
        break;
    }

    input.popLimit(oldLimit);

    return true;
  }

  private void handleMessage(Proto.SessionHeader header) {
    // Nothing to do.
  }

  private void handleMessage(Proto.CallRequest request,
                             ExecutorService executor) throws IOException {
    if (request.getPrerequisiteCallIdCount() > 0) {
      // TODO(kenton):  Implement prerequisites.
      sendCallError(request.getCallId(),
                    Proto.CallError.Type.PREREQUISITE_FAILED,
                    "Prerequisites not implemented.");
      return;
    }

    int serviceId = request.getTargetServiceId();
    if (serviceId < 0 || serviceId >= channels.size()) {
      sendCallError(request.getCallId(),
                    Proto.CallError.Type.BAD_REQUEST,
                    "Invalid target service ID.");
      return;
    }

    Channel channel = channels.get(serviceId);

    Proto.CallPayload payload = request.getPayload();

    Channel[] receivedServices =
      new Channel[payload.getEmbeddedServiceIdCount()];
    for (int i = 0; i < receivedServices.length; i++) {
      receivedServices[i] = new ChannelImpl(payload.getEmbeddedServiceId(i));
    }

    RpcContext controller =
        new RpcContext(Arrays.asList(receivedServices));

    // Run in another thread so that we can continue handling requests in this
    // one.
    executor.submit(new CallTask(request, channel, controller));
  }

  private class CallTask implements Runnable {
    Proto.CallRequest request;
    Channel channel;
    RpcContext controller;

    CallTask(Proto.CallRequest request, Channel channel,
             RpcContext controller) {
      this.request = request;
      this.channel = channel;
      this.controller = controller;
    }

    public void run() {
      ByteString applicationResponse;

      try {
        try {
          applicationResponse = channel.callMethod(
              request.getMethodName(), controller,
              request.getPayload().getData());
        } catch (ServiceException e) {
          sendCallError(request.getCallId(),
                        Proto.CallError.Type.APPLICATION_ERROR,
                        e.getMessage());
          return;
        } catch (Throwable e) {
          e.printStackTrace();
          sendCallError(request.getCallId(),
                        Proto.CallError.Type.SERVER_ERROR,
                        "Uncaught exception: " + e.getMessage());
          return;
        }

        writeMessage(Proto.Stream.CALL_RESPONSE_FIELD_NUMBER,
          Proto.CallResponse.newBuilder()
            .setCallId(request.getCallId())
            .setPayload(
                makePayload(applicationResponse,
                        controller.getResponseServiceReceiver().receiveAll()))
            .build());
      } catch (IOException e) {
        // TODO(kenton):  The output stream seems to be broken.  What can we do
        //   here?
        e.printStackTrace();
      }
    }
  }

  private Proto.CallPayload makePayload(
      ByteString applicationMessage,
      List<Channel> list) {
    Proto.CallPayload.Builder responsePayloadBuilder =
      Proto.CallPayload.newBuilder();

    responsePayloadBuilder.setData(applicationMessage);

    for (Channel sentService: list) {
      responsePayloadBuilder.addEmbeddedServiceId(channels.size());
      channels.add(sentService);
    }

    return responsePayloadBuilder.build();
  }

  private void sendCallError(int callId, Proto.CallError.Type type,
                             String description) throws IOException {
    writeMessage(Proto.Stream.CALL_RESPONSE_FIELD_NUMBER,
      Proto.CallResponse.newBuilder()
        .setCallId(callId)
        .setError(Proto.CallError.newBuilder()
            .setType(type).setDescription(description))
        .build());
  }

  private void handleMessage(Proto.CallResponse response) throws IOException {
    DoneCallback callback;
    synchronized (pendingCalls) {
      callback = pendingCalls.get(response.getCallId());
    }
    if (callback == null) {
      sendError(Proto.MiscError.Type.PROTOCOL_ERROR,
                "Got CallResponse with invalid call ID.");
      throw new ProtocolErrorException(
          "Got CallResponse with invalid call ID.");
    } else {
      callback.run(response);
    }
  }

  private void handleMessage(Proto.CancelCall request) {
    // TODO:  Implement cancellation.
  }

  private void handleMessage(Proto.MiscError error) {
    System.err.println("RPC error: " + error.getDescription());
  }

  private void sendError(Proto.MiscError.Type type,
                         String description)
      throws IOException {
    writeMessage(Proto.Stream.ERROR_FIELD_NUMBER,
      Proto.MiscError.newBuilder()
        .setType(type)
        .setDescription(description)
        .build());
  }

  private void sendError(Proto.MiscError.Type type,
                         String description, int number)
      throws IOException {
    writeMessage(Proto.Stream.ERROR_FIELD_NUMBER,
      Proto.MiscError.newBuilder()
        .setType(type)
        .setDescription(description)
        .setUnknownFieldNumber(number)
        .build());
  }

  private void writeMessage(int fieldNumber, MessageLite message)
      throws IOException {
    synchronized (output) {
      output.writeMessage(fieldNumber, message);
      output.flush();
    }
  }

  private static class DoneCallback {
    Proto.CallResponse response = null;

    public synchronized void run(Proto.CallResponse response) {
      this.response = response;
      this.notify();
    }

    public synchronized Proto.CallResponse waitUntilCalled()
        throws InterruptedException {
      while (response == null) {
        wait();
      }
      return response;
    }
  }

  private class ChannelImpl implements Channel {
    private final int serviceId;

    ChannelImpl(int serviceId) {
      this.serviceId = serviceId;
    }

    @SuppressWarnings("unchecked")
    @Override
    public <ResponseType extends MessageLite> ResponseType callMethod(
        String methodName, RpcContext context,
        MessageLite request, ResponseType responsePrototype)
        throws ServiceException {
      ByteString response =
          callMethod(methodName, context, request.toByteString());

      try {
        return (ResponseType) responsePrototype.newBuilderForType()
            .mergeFrom(response).build();
      } catch (UninitializedMessageException e) {
        throw new ServiceException(
            "Could not parse response message: " + e.getMessage(),
            Proto.CallError.Type.SERVER_ERROR, e);
      } catch (InvalidProtocolBufferException e) {
        throw new ServiceException(
            "Could not parse response message: " + e.getMessage(),
            Proto.CallError.Type.SERVER_ERROR, e);
      }
    }

    @Override
    public ByteString callMethod(String methodName,
        RpcContext context, ByteString request) throws ServiceException {
      int callId = outgoingCallIdCounter.incrementAndGet();

      DoneCallback done = new DoneCallback();
      Proto.CallResponse response;

      synchronized (pendingCalls) {
        pendingCalls.put(callId, done);
      }

      try {
        try {
          writeMessage(Proto.Stream.CALL_REQUEST_FIELD_NUMBER,
            Proto.CallRequest.newBuilder()
              .setCallId(callId)
              .setTargetServiceId(serviceId)
              .setMethodName(methodName)
              .setPayload(makePayload(request,
                  context.getRequestServiceReceiver().receiveAll()))
              .build());
        } catch (IOException e) {
          ServiceException e2 =
            new ServiceException("I/O error when sending RPC request.");
          e2.initCause(e);
          throw e2;
        }

        try {
          response = done.waitUntilCalled();
        } catch (InterruptedException e) {
          ServiceException e2 = new ServiceException(
              "Caught InterruptedException while waiting for RPC.");
          e2.initCause(e);
          throw e2;
        }
      } finally {
        synchronized (pendingCalls) {
          pendingCalls.remove(callId);
        }
      }

      if (response.hasError()) {
        // TODO(kenton):  Decide when to retry.
        throw new ServiceException(response.getError().getDescription());
      }

      Proto.CallPayload payload = response.getPayload();

      Channel[] receivedServices =
        new Channel[payload.getEmbeddedServiceIdCount()];
      for (int i = 0; i < receivedServices.length; i++) {
        receivedServices[i] = new ChannelImpl(payload.getEmbeddedServiceId(i));
      }

      context.setResponseServices(Arrays.asList(receivedServices));

      return payload.getData();
    }
  }

  private static class NullChannel implements Channel {
    @Override
    public ByteString callMethod(String methodName,
        RpcContext context, ByteString request) throws ServiceException {
      throw new ServiceException(
          "This connection does not have a default service.",
          Proto.CallError.Type.BAD_REQUEST);
    }

    @Override
    public <ResponseType extends MessageLite> ResponseType callMethod(
        String methodName, RpcContext context,
        MessageLite request, ResponseType responsePrototype)
        throws ServiceException {
      throw new ServiceException(
          "This connection does not have a default service.",
          Proto.CallError.Type.BAD_REQUEST);
    }
  }
}
